/*
 *  Copyright (C) 2025 ameise <ameise.wang@gmail.com> - All Rights Reserved
 *
 *  This file is part of e3utils.
 *
 *  e3utils is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  e3utils is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with e3utils. If not, see <https://www.gnu.org/licenses/>.
 */

package load_balancer

import (
	"gitee.com/ameise84/e3ds/rbmap"
)

func NewBalancerHash[T Context]() Balancer[T] {
	mp := rbmap.New[uint64, T]()
	return &hashBalancer[T]{
		mp:   mp,
		iter: mp.BeginIterator(),
	}
}

type hashBalancer[T Context] struct {
	mp    rbmap.RBMap[uint64, T]
	iter  *rbmap.Iterator[uint64, T]
	zeroT T
}

func (b *hashBalancer[T]) Mode() Mode {
	return HashMode
}

func (b *hashBalancer[T]) Count() int {
	return b.mp.Size()
}

func (b *hashBalancer[T]) Register(ctx T) error {
	id := ctx.LoadBalancerID()
	if _, ok := b.mp.Load(id); ok {
		return ErrIdRepeat
	}
	b.mp.Store(id, ctx)
	return nil
}

func (b *hashBalancer[T]) UnRegister(id uint64) (v T, isFind bool) {
	key, err := b.iter.Key()
	if err == nil && key == id {
		val := b.iter.ValueNoError()
		b.iter, _ = b.iter.Delete()
		v, isFind = val, true
	} else {
		if val, ok := b.mp.Delete(id); ok {
			v, isFind = val, true
		}
	}
	return
}

func (b *hashBalancer[T]) Clean() []T {
	out := make([]T, 0, b.mp.Size())
	mp := b.mp.Clean()
	for _, v := range mp {
		out = append(out, v)
	}
	return out
}

func (b *hashBalancer[T]) Range(f func(T) bool) {
	b.mp.Range(func(iter *rbmap.Iterator[uint64, T]) bool {
		return f(iter.ValueNoError())
	})
}

func (b *hashBalancer[T]) TakeByID(id uint64) (T, bool) {
	iter := b.mp.Search(id)
	if iter == b.mp.EndIterator() {
		return b.zeroT, false
	}
	return iter.ValueNoError(), true
}

func (b *hashBalancer[T]) TakeNext(id uint64) (T, error) {
	iter := b.mp.Search(id, rbmap.SearchModeGT|rbmap.SearchModeET)
	if iter == b.mp.EndIterator() {
		if b.mp.Size() > 0 {
			return b.mp.BeginIterator().ValueNoError(), nil
		}
		return b.zeroT, ErrIdNotFind
	}
	return iter.ValueNoError(), nil
}

func (b *hashBalancer[T]) SetAutoPriWt(int16) {}

func (b *hashBalancer[T]) UpdatePriWt(uint64, int) {}
